package tunnel

import (
	"bytes"
	"encoding/binary"
	"log"
	"net"
	"sync"
)

const (
	BUFF_SIZE                    = 1024 * 20
	TOKEN_SERVER_CONNECTED uint8 = 200 + iota
	TOKEN_SERVER_STATUS
	TOKEN_SERVER_DISCONNECT
	TOKEN_SERVER_FORWARD
	TOKEN_SERVER_REFORWARD
	TOKEN_SERVER_DISCONNECT_ALL
)

var (
	buffer = sync.Pool{
		New: func() interface{} { return make([]byte, BUFF_SIZE) },
	}
	buffer_new = sync.Pool{
		New: func() interface{} { return new(bytes.Buffer) },
	}
)

type Tunnel_server struct {
	tunnel_local_clients_map  map[uint16]net.Conn
	remote_listener           net.Listener
	local_listener            net.Listener
	tunnel_bind_remote_client net.Conn
	locker                    sync.RWMutex
	tunnel_session_id         uint16
	local_bind_addr           string
	remote_bind_addr          string
	is_running                bool
}

func pack(conn net.Conn, token uint8, session_id uint16, data []byte) {
	var buf = buffer_new.Get().(*bytes.Buffer)
	buf.WriteByte(token)
	binary.Write(buf, binary.LittleEndian, session_id)
	buf.Write(data)
	if conn != nil {
		buf.WriteTo(conn)
	}
}
func New_Tunnel_server(remote_addr string, local_addr string) *Tunnel_server {
	server := new(Tunnel_server)
	server.tunnel_local_clients_map = make(map[uint16]net.Conn, 20)
	server.remote_bind_addr = remote_addr
	server.local_bind_addr = local_addr
	server.is_running = true
	return server
}
func (self *Tunnel_server) getLocalBindConnSessionId(c net.Conn) uint16 {
	for k, v := range self.tunnel_local_clients_map {
		if c != nil && v != nil && c == v {
			return k
		}
	}
	return 0
}
func (self *Tunnel_server) getLocalBindSessionConn(session_id uint16) net.Conn {
	return self.tunnel_local_clients_map[session_id]
}
func (self *Tunnel_server) clearLocalBindClients() {
		log.Println("关闭本地所有连结")
	for k, v := range self.tunnel_local_clients_map {
		v.Close()
		delete(self.tunnel_local_clients_map, k)
	}

}
func (self *Tunnel_server) onTunnelLocalBindConnect(c net.Conn) {
	defer c.Close()
	self.locker.Lock()
	self.tunnel_local_clients_map[self.tunnel_session_id] = c
	self.tunnel_session_id += 1
	self.locker.Unlock()

	for {
		data := buffer.Get().([]byte)
		n, err := c.Read(data)
		if err != nil {
			self.locker.Lock()
			delete(self.tunnel_local_clients_map, self.getLocalBindConnSessionId(c))
			self.locker.Unlock()
			break
		}
		session_id := self.getLocalBindConnSessionId(c)
		if self.tunnel_bind_remote_client != nil && session_id > 0 {
			log.Println("收到数据", session_id, n)
			log.Println("转发数据", session_id, n)
			pack(self.tunnel_bind_remote_client, TOKEN_SERVER_FORWARD, session_id, data[:n])
		}

	}
}
func (self *Tunnel_server) onTunnelRemoteBindConnect(c net.Conn) {
	defer c.Close()
	c.(*net.TCPConn).SetKeepAlive(true)
	if self.tunnel_bind_remote_client != nil {
		self.tunnel_bind_remote_client.Close()
	}
	self.tunnel_bind_remote_client = c

	for {
		data := buffer.Get().([]byte)
		n, err := c.Read(data)
		if err != nil {
			self.tunnel_bind_remote_client = nil
			break
		}
		session_id := binary.LittleEndian.Uint16(data[1:3])
		conn := self.getLocalBindSessionConn(session_id)
		switch data[0] {
		case TOKEN_SERVER_DISCONNECT:
			if conn != nil {
				conn.Close()
			}
		case TOKEN_SERVER_DISCONNECT_ALL:
			self.clearLocalBindClients()
		case TOKEN_SERVER_STATUS:
			pack(self.tunnel_bind_remote_client, TOKEN_SERVER_DISCONNECT, session_id, []byte{})
		case TOKEN_SERVER_REFORWARD:
			log.Println("收到转发数据回复", session_id, n)
			if conn != nil {
				conn.Write(data[3:n])
			}
		}

	}
}
func (self *Tunnel_server) runServer(err error, addr string, onConnect func(c net.Conn)) {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		log.Println(err)
		return
	}

	if addr == self.local_bind_addr {
		log.Println("local bind", listener.Addr().String(), "running!")
		self.local_listener = listener
	} else if addr == self.remote_bind_addr {
		log.Println("remote bind", listener.Addr().String(), "running!")
		self.remote_listener = listener
	}
	defer listener.Close()
	for self.is_running {
		conn, err := listener.Accept()
		if err != nil {
			continue
		}
		go onConnect(conn)
	}
}
func (self *Tunnel_server) Start() error {
	var err error
	go self.runServer(err, self.remote_bind_addr, self.onTunnelRemoteBindConnect)
	if err != nil {
		return err
	}
	go self.runServer(err, self.local_bind_addr, self.onTunnelLocalBindConnect)
	if err != nil {
		self.is_running = false
		return err
	}
	return err
}
func (self *Tunnel_server) stop() {
	self.is_running = false
	self.local_listener.Close()
	self.remote_listener.Close()
}
